"""
GPU 使用示例
功能：演示如何在 Modal 中请求和使用 GPU 资源
特点：支持单个和多个 GPU 配置（A10G、H100 等）
运行：modal run 05_gpu_usage.py
"""
import modal

app = modal.App("gpu-app")
image = modal.Image.debian_slim().pip_install("torch")

# 请求一个 NVIDIA A10G GPU
@app.function(image=image, gpu="A10G")
def use_gpu():
    import torch
    if torch.cuda.is_available():
        print("GPU is available:", torch.cuda.get_device_name(0))
    else:
        print("GPU not found!")

# 请求 4 个 H100 GPU
@app.function(image=image, gpu="H100:4")
def use_multi_gpu():
    import torch
    print(f"Found {torch.cuda.device_count()} GPUs.")

@app.local_entrypoint()
def main():
    use_gpu.remote()
    use_multi_gpu.remote()